#include "PotFitter.h"


void PotFitter::Init(vector<string> &ar){
	cout << "==========================================\n";	
	MAInput1 inp(ar);
	
	cout<<" Loading force field ..."<<endl;
	string sForceField="OSS2";	inp.getString("-ff",sForceField);		
	if((sForceField=="OSS2")||(sForceField=="OSS2_ITER")){	
		of=new PotWater();				
	}else{
		if(sForceField=="HF_OSS2")	of=new PotHF();
		
	}		
	assert(of);
	
	of->SetModel_(sForceField.c_str());	
	int unit; if(inp.getInt("-u",unit)) of->SetUnit_(unit);

	int vi;
	if(inp.getInt("-miter",vi)) of->SetMaxEval_(vi);
	if(inp.getInt("-ls",vi)) of->SetOptMethod_(vi);
	double vd; 
	if(inp.getDouble("-mstep",vd)) of->SetMaxStep_(vd);
	if(inp.getDouble("-ftol",vd)) of->SetFcnTol_(vd);
	if(inp.getDouble("-gtol",vd)) of->SetGradTol_(vd);
	
	of->PrintInfo();
	
	cout << "Loading bounds..." << endl;
	string boundfile; assert(inp.getString("-a",boundfile));
	ifstream fin(boundfile.c_str()); assert(fin);

	nparam_inp=nfixed=0;	label_inp.clear();
	
	int i=0;
	while(!fin.eof()){
		string s=""; 	getline(fin,s);		if(s=="") break;
		
		stringstream ss; ss<<s<<" ";
		
		string tName;
		
		ss>>param_inp[i]>>lb_inp[i]>>ub_inp[i]>>fixed[i]>>tName;
		label_inp.push_back(tName);
		//sscanf(s.c_str(),"%lf %lf %lf %d %s", &param_inp[nparam_inp], &lb_inp[nparam_inp], &ub_inp[nparam_inp], &fixed[nparam_inp]);
		
		nfixed+=fixed[i];
		
		//cout<<i << ":  " << param_inp[i] << " " << lb_inp[i] << " " << ub_inp[i] << " " <<  fixed[i] << " "<<label_inp[i]<<endl;	
		
		i++;
	}
	fin.close();
	
	nparam_inp=i;

	nparam = nparam_inp - nfixed;
	
	label.clear();
	
	int t = 0;
	for(int i=0; i<nparam_inp; i++){
		if (!fixed[i]){
			lb[t] = lb_inp[i];			ub[t] = ub_inp[i];		
			param[t]=param_inp[i];	
			label.push_back(label_inp[i]);	
			assert(lb[t]<=ub[t]);
			t++;
		}
	}
	

	cout << nparam << " variables, " << nfixed << " fixed values...done." << endl;
	
	printParam(cout);

    cout << "\n============================  Reference structures...\n";
	
	string reffile; assert(inp.getString("-ref",reffile));
	
	ref_inp.clear();	loadData(reffile.c_str(),F_XYZ,ref_inp);

	for(int i=0;i<ref_inp.size();i++){	ref_inp[i].Write(cout,F_XYZ);}

	cout << "\n=============================  Data structures...\n";
	
		
	string s,datafile;
	
	data.clear();
	for(int i=0;i<50;i++)
		if(inp.getString(str(format("-i%d")%(i+1)),datafile)){
			int pos1=data.size();
			loadData(datafile.c_str(),F_XYZ,data); assert(data.size()<MAX_DATA_SIZE);
			double w=1;
			inp.getDouble(str(format("-w%d")%(i+1)),w);
			
			for(int j=pos1;j<data.size();j++)	weight[j]=w;
				
		}
	
	cout << "Compute observation data ...";
	
	ProcessData();
	
	cout << " ... " << n_measure << " observations ." << endl;
	
	cout << "==========================================\n";	
	
}

int PotFitter::loadData(const char* filename,int inmode,vector<Individual> &d){
	ifstream fi(filename); assert(fi);	
    cout << "Loading data from " << filename ;
	
	Individual inp;
	int i=0;
	while (!fi.eof()) {		
		if(!inp.Read(fi,inmode)) break;
		inp.CorrectOrder();
		d.push_back(inp);		
		i++;
	}

	cout << " ... " << i  << " inputs..." << endl;
	fi.close();
	return i;
}

void PotFitter::ProcessData(){
	n_measure=0;
	
	for (int i=0; i< data.size(); i++){
		int nF=data[i].nAtom - data[i].nType[HYDROGEN];
				
		measure[n_measure]= data[i].energy - (nF-1)*ref_inp[0].energy ;
		switch((int)(data[i].TotalCharge())){
			case 0:  measure[n_measure]-= ref_inp[0].energy; break;
			case 1:  measure[n_measure]-= ref_inp[1].energy; break;
			case -1:  measure[n_measure]-= ref_inp[2].energy; break;
		}
		measure[n_measure]*=sqrt(weight[i]);
		
		n_measure++;
	}
	/*
	for (int i=0; i< data.size(); i++){		
		measure[n_measure]=data[i].rmsGrad;	
		n_measure++;	
	}
	*/	
}

void PotFitter::levmar_func(double *p, double *y, int m, int n, void *d){
	PotFitter* t=(PotFitter*)d;		
	//checking parameters
	assert(m==t->nparam);		assert(n==t->n_measure);
	t->evaluate(p,y,0);
}

double PotFitter::evaluate(double *p,double *y,int debug){	
	assert(of);
	double alpha[100];
	ConvertParam(p,alpha);
	of->ReadParam(alpha);
	
	if(debug>=2){
		cout << "==========================================\n";
		cout << "@# Alpha: ";
		for(int i=0; i<nparam_inp; i++)	printf("%.10lf ", alpha[i]);
		cout << endl;
	}
		
	//calculate 
	double ener_shift[3];	
	for(int i=0; i<3; i++){	
		of->SetDim_(ref_inp[i].nAtom*DGR);
		ener_shift[i] = of->evaluate_(ref_inp[i].x);		
		assert(ener_shift[i]==ener_shift[i]);
	}	

	if (debug>=3){
		 cout << "Energy shift values: neu= "<< ener_shift[0] << " pron= " << ener_shift[1] << " dep= " << ener_shift[2] << endl;
		 cout <<"\n\n@$ inputs# nAtoms Weight Obsr.Data Calc.Data deltaE" << endl;
	}
	
	double rms=0;
	
	int c=0;
	for (int i=0; i<data.size(); i++){
		//calculate binding energy
		of->SetDim_(data[i].nAtom*3);
		
		int nF=data[i].nAtom - data[i].nType[HYDROGEN];
		
		of->SetExtra_(&nF);
		
		//of->PrintInfo();
		
		double ener = of->evaluate_(data[i].x);	
		
		y[c] = ener - (nF-1)*ener_shift[0] ;
		switch((int)(data[i].TotalCharge())){
			case 0:  y[c]-= ener_shift[0]; break;
			case 1:  y[c]-= ener_shift[1]; break;
			case -1:  y[c]-= ener_shift[2]; break;
		}	
		y[c]*=sqrt(weight[c]);
		
		double dE = fabs(y[c] - measure[c]);		rms+=dE*dE;
		
				
		if(debug>=4) printf("@$\t%5d\t%5d\t%6.2lf\t%12.8lf\t%12.8lf\t%12.8lf\n", c, data[i].nAtom, weight[i], measure[c], y[c] ,dE);
		
		c++;
	}
		
	/*
	double g[SIZE_MAX*DGR];
	for (int i=0; i<data.size(); i++){		
			//calculate force
		of->SetDim_(data[i].nAtom*3);
		double ener = of->gradient_num(data[i].x,g);
		y[c] = ener;		
		double dE = fabs(y[c] - measure[i]);		rms+=dE*dE;
		if (0) printf("@$ %5d   %5d  %12.8lf %12.8lf %12.8lf\n", c, data[i].nAtom, y[c], measure[c] ,dE);			
	c++;		
	}	*/
	
	rms=sqrt(rms/n_measure);
	
	if (debug>=1) cout<<"@$ Total RMS = "<<rms<<" \n\n";
	
	return rms;
}

void PotFitter::ConvertParam(double *actual,double *p){	
	int c=0;
	for(int i=0; i<nparam_inp; i++) 
		if (!fixed[i]){
			p[i] = actual[c];
			c++;
		}else p[i] = param_inp[i];
}

void PotFitter::printParam(ostream &os,double* x){
	double res[100];

	if(x)	ConvertParam(x,res);
	else	ConvertParam(param,res);
	
	for(int i=0; i<nparam_inp; i++) 
		os<<format("%15.8lf\t%15.8lf\t%15.8lf\t%d")%res[i]% lb_inp[i]% ub_inp[i]% fixed[i]<<" "<< label_inp[i]<<endl;	
}

int PotFitter::RunLevMar(int nRuns,bool bBound,int debug){
	//if(nRuns==0) return 0;
	assert(of);	
	
	int m, n;  m=getNParam(); n=getNMeasure();
	
	double opts[LM_OPTS_SZ]; //options for reparameterizing
	opts[0]=LM_INIT_MU; opts[1]=1E-15; opts[2]=1E-15; opts[3]=1E-20;
	opts[4]=LM_DIFF_DELTA; // relevant only if the finite difference Jacobian version is used 
	
	double info[LM_INFO_SZ]; //info after reparameterizing
	
	double *work=new double[(LM_DIF_WORKSZ(m, n)+m*m)]; assert(work);
	double*	covar=work+LM_DIF_WORKSZ(m, n);
	
	
	if(debug){
		cout<<" \n=========== Start running =========\n";
		cout <<" Using "<<of->equation()<<" and "<<bBound<<" boundary "<<endl;
	}
	int ret;
	if(!bBound)
		ret=dlevmar_dif(PotFitter::levmar_func, getParam(), getMeasure(), m, n, nRuns, opts, info, work, covar, this);
	else
		ret=dlevmar_bc_dif(PotFitter::levmar_func, getParam(), getMeasure(), m, n, lb, ub, nRuns, opts, info, work, covar, this);
	
	delete[] work;
	
	

	if(debug){		
		cout<<" \n=========== Stop running =========\n";
		cout<<"\n Covariance of the fit:\n";
		cout<<"Param ";
		for(int i=0; i<m; ++i) cout<<label[i]<<" ";
		cout<<endl;
		for(int i=0; i<m; ++i){
			cout<<label[i]<<" ";
			for(int j=0;j<m;j++)			
				cout<<format("%1.2g ")%covar[i*m+j];
			cout<<endl;
		}		
	
		
	
	cout<<"\n\n Levenberg-Marquardt's Minimization info : ";	
		
	switch((int)info[6]){ //=reason for terminating: 
		case 1: cout<<" - stopped by small gradient J^T e \n"; break;
		case 2: cout<<" - stopped by small Dp \n"; break;
		case 3: cout<<" - stopped by itmax \n"; break;
		case 4: cout<<" - singular matrix. Restart from current p with increased mu \n"; break;
		case 5: cout<<" - no further error reduction is possible. Restart with increased mu \n"; break;
		case 6: cout<<" - stopped by small ||e||_2 \n"; break;
		case 7: cout<<" - stopped by invalid (i.e. NaN or Inf) func values. This is a user error \n"; break;
	}
	
	cout<<info[0]<<" = ||e||_2 at initial p. \n";
	cout<<info[1]<<" = ||e||_2 at estimated p. \n";
	cout<<info[2]<<" = ||J^T e||_inf \n";
	cout<<info[3]<<" = ||Dp||_2  \n";
	cout<<info[4]<<" = mu/max[J^T J]_ii ]  \n";	
	
	cout<<(int)info[5]<<"= # iterations, \n";
	cout<<(int)info[7]<<" = # function evaluations \n";
	cout<<(int)info[8]<<" = # Jacobian evaluations \n";		
	}

	return info[5];
}
